import unittest
import os
import sys
import pathlib
import urllib
import shutil
import re

import numpy as np
import torch
import matplotlib.pyplot as plt

%load_ext autoreload
%autoreload 2

test = unittest.TestCase()
plt.rcParams.update({'font.size': 12})
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)

Using device: cuda

CORPUS_URL = 'https://github.com/cedricdeboom/character-level-rnn-datasets/raw/master/datasets/shakespeare.txt'
DATA_DIR = pathlib.Path.home().joinpath('.pytorch-datasets')

def download_corpus(out_path=DATA_DIR, url=CORPUS_URL, force=False):
    pathlib.Path(out_path).mkdir(exist_ok=True)
    out_filename = os.path.join(out_path, os.path.basename(url))
    
    if os.path.isfile(out_filename) and not force:
        print(f'Corpus file {out_filename} exists, skipping download.')
    else:
        print(f'Downloading {url}...')
        with urllib.request.urlopen(url) as response, open(out_filename, 'wb') as out_file:
            shutil.copyfileobj(response, out_file)
        print(f'Saved to {out_filename}.')
    return out_filename
    
corpus_path = download_corpus()

Corpus file /home/eran.b/.pytorch-datasets/shakespeare.txt exists, skipping download.

with open(corpus_path, 'r', encoding='utf-8') as f:
    corpus = f.read()

print(f'Corpus length: {len(corpus)} chars')
print(corpus[7:1234])

Corpus length: 6347703 chars
ALLS WELL THAT ENDS WELL

by William Shakespeare

Dramatis Personae

  KING OF FRANCE
  THE DUKE OF FLORENCE
  BERTRAM, Count of Rousillon
  LAFEU, an old lord
  PAROLLES, a follower of Bertram
  TWO FRENCH LORDS, serving with Bertram

  STEWARD, Servant to the Countess of Rousillon
  LAVACHE, a clown and Servant to the Countess of Rousillon
  A PAGE, Servant to the Countess of Rousillon

  COUNTESS OF ROUSILLON, mother to Bertram
  HELENA, a gentlewoman protected by the Countess
  A WIDOW OF FLORENCE.
  DIANA, daughter to the Widow

  VIOLENTA, neighbour and friend to the Widow
  MARIANA, neighbour and friend to the Widow

  Lords, Officers, Soldiers, etc., French and Florentine  

SCENE:
Rousillon; Paris; Florence; Marseilles

ACT I. SCENE 1.
Rousillon. The COUNT'S palace

Enter BERTRAM, the COUNTESS OF ROUSILLON, HELENA, and LAFEU, all in black

  COUNTESS. In delivering my son from me, I bury a second husband.
  BERTRAM. And I in going, madam, weep o'er my father's death anew;
    but I must attend his Majesty's command, to whom I am now in
    ward, evermore in subjection.
  LAFEU. You shall find of the King a husband, madam; you, sir, a
    father. He that so generally is at all times good must of

import hw3.charnn as charnn

char_to_idx, idx_to_char = charnn.char_maps(corpus)
print(char_to_idx)

test.assertEqual(len(char_to_idx), len(idx_to_char))
test.assertSequenceEqual(list(char_to_idx.keys()), list(idx_to_char.values()))
test.assertSequenceEqual(list(char_to_idx.values()), list(idx_to_char.keys()))

{'\n': 0, ' ': 1, '!': 2, '"': 3, '$': 4, '&': 5, "'": 6, '(': 7, ')': 8, ',': 9, '-': 10, '.': 11, '0': 12, '1': 13, '2': 14, '3': 15, '4': 16, '5': 17, '6': 18, '7': 19, '8': 20, '9': 21, ':': 22, ';': 23, '<': 24, '?': 25, 'A': 26, 'B': 27, 'C': 28, 'D': 29, 'E': 30, 'F': 31, 'G': 32, 'H': 33, 'I': 34, 'J': 35, 'K': 36, 'L': 37, 'M': 38, 'N': 39, 'O': 40, 'P': 41, 'Q': 42, 'R': 43, 'S': 44, 'T': 45, 'U': 46, 'V': 47, 'W': 48, 'X': 49, 'Y': 50, 'Z': 51, '[': 52, ']': 53, '_': 54, 'a': 55, 'b': 56, 'c': 57, 'd': 58, 'e': 59, 'f': 60, 'g': 61, 'h': 62, 'i': 63, 'j': 64, 'k': 65, 'l': 66, 'm': 67, 'n': 68, 'o': 69, 'p': 70, 'q': 71, 'r': 72, 's': 73, 't': 74, 'u': 75, 'v': 76, 'w': 77, 'x': 78, 'y': 79, 'z': 80, '}': 81, '\ufeff': 82}

corpus, n_removed = charnn.remove_chars(corpus, ['}','$','_','<','\ufeff'])
print(f'Removed {n_removed} chars')

# After removing the chars, re-create the mappings
char_to_idx, idx_to_char = charnn.char_maps(corpus)

Removed 34 chars

# Wrap the actual embedding functions for calling convenience
def embed(text):
    return charnn.chars_to_onehot(text, char_to_idx)

def unembed(embedding):
    return charnn.onehot_to_chars(embedding, idx_to_char)

text_snippet = corpus[3104:3148]
print(text_snippet)
print(embed(text_snippet[0:3]))

test.assertEqual(text_snippet, unembed(embed(text_snippet)))
test.assertEqual(embed(text_snippet).dtype, torch.int8)

brine a maiden can season her praise in.
   
tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
         0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0]], dtype=torch.int8)

# Create dataset of sequences
seq_len = 64
vocab_len = len(char_to_idx)

# Create labelled samples
samples, labels = charnn.chars_to_labelled_samples(corpus, char_to_idx, seq_len, device)
print(f'samples shape: {samples.shape}')
print(f'labels shape: {labels.shape}')

# Test shapes
num_samples = (len(corpus) - 1) // seq_len
test.assertEqual(samples.shape, (num_samples, seq_len, vocab_len))
test.assertEqual(labels.shape, (num_samples, seq_len))

# Test content
for _ in range(1000):
    # random sample
    i = np.random.randint(num_samples, size=(1,))[0]
    # Compare to corpus
    test.assertEqual(unembed(samples[i]), corpus[i*seq_len:(i+1)*seq_len], msg=f"content mismatch in sample {i}")
    # Compare to labels
    sample_text = unembed(samples[i])
    label_text = str.join('', [idx_to_char[j.item()] for j in labels[i]])
    test.assertEqual(sample_text[1:], label_text[0:-1], msg=f"label mismatch in sample {i}")

samples shape: torch.Size([99182, 64, 78])
labels shape: torch.Size([99182, 64])

import re
import random

i = random.randrange(num_samples-5)
for i in range(i, i+5):
    test.assertEqual(len(samples[i]), seq_len)
    s = re.sub(r'\s+', ' ', unembed(samples[i])).strip()
    print(f'sample [{i}]:\n\t{s}')

sample [93594]:
	I prize it As I weigh grief, which I would spare: for honour, 'T
sample [93595]:
	is a derivative from me to mine, And only that I stand for. I ap
sample [93596]:
	peal To your own conscience, sir, before Polixenes Came to your
sample [93597]:
	court, how I was in your grace, How merited to be so; since he c
sample [93598]:
	ame, With what encounter so uncurrent I Have strain'd to appear

from hw3.charnn import SequenceBatchSampler

sampler = SequenceBatchSampler(dataset=range(32), batch_size=10)
sampler_idx = list(sampler)
print('sampler_idx =\n', sampler_idx)

# Test the Sampler
test.assertEqual(len(sampler_idx), 30)
batch_idx = np.array(sampler_idx).reshape(-1, 10)
for k in range(10):
    test.assertEqual(np.diff(batch_idx[:, k], n=2).item(), 0)

sampler_idx =
 [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29]

/home/eran.b/miniconda3/envs/cs236781-hw/lib/python3.8/site-packages/torch/utils/data/sampler.py:65: UserWarning: `data_source` argument is not used and will be removed in 2.2.0.You may still have custom implementation that utilizes it.
  warnings.warn("`data_source` argument is not used and will be removed in 2.2.0."

import torch.utils.data

# Create DataLoader returning batches of samples.
batch_size = 32

ds_corpus = torch.utils.data.TensorDataset(samples, labels)
sampler_corpus = SequenceBatchSampler(ds_corpus, batch_size)
dl_corpus = torch.utils.data.DataLoader(ds_corpus, batch_size=batch_size, sampler=sampler_corpus, shuffle=False)

print(f'num batches: {len(dl_corpus)}')

x0, y0 = next(iter(dl_corpus))
print(f'shape of a batch of samples: {x0.shape}')
print(f'shape of a batch of labels: {y0.shape}')

num batches: 3100
shape of a batch of samples: torch.Size([32, 64, 78])
shape of a batch of labels: torch.Size([32, 64])

# Check that sentences in in same index of different batches complete each other.
k = random.randrange(batch_size)
for j, (X, y) in enumerate(dl_corpus,):
    print(f'=== batch {j}, sample {k} ({X[k].shape}): ===')
    s = re.sub(r'\s+', ' ', unembed(X[k])).strip()
    print(f'\t{s}')
    if j==4: break

=== batch 0, sample 24 (torch.Size([64, 78])): ===
	mind were clear again, that I might water an ass at it. I ha
=== batch 1, sample 24 (torch.Size([64, 78])): ===
	d rather be a tick in a sheep than such a valiant ignorance.
=== batch 2, sample 24 (torch.Size([64, 78])): ===
	Exit ACT IV. SCENE 1. Troy.
=== batch 3, sample 24 (torch.Size([64, 78])): ===
	A street Enter, at one side, AENEAS, and servant with a torch;
=== batch 4, sample 24 (torch.Size([64, 78])): ===
	at another, PARIS, DEIPHOBUS, ANTENOR, DIOMEDES the Grecian, an

in_dim = vocab_len
h_dim = 256
n_layers = 3
model = charnn.MultilayerGRU(in_dim, h_dim, out_dim=in_dim, n_layers=n_layers)
model = model.to(device)
print(model)

# Test forward pass
y, h = model(x0.to(dtype=torch.float, device=device))
print(f'y.shape={y.shape}')
print(f'h.shape={h.shape}')

test.assertEqual(y.shape, (batch_size, seq_len, vocab_len))
test.assertEqual(h.shape, (batch_size, n_layers, h_dim))
test.assertEqual(len(list(model.parameters())), 9 * n_layers + 2)

MultilayerGRU(
  (layer_params): ModuleList(
    (0): ModuleList(
      (0): Linear(in_features=78, out_features=256, bias=True)
      (1): Linear(in_features=256, out_features=256, bias=False)
      (2): Linear(in_features=78, out_features=256, bias=True)
      (3): Linear(in_features=256, out_features=256, bias=False)
      (4): Linear(in_features=78, out_features=256, bias=True)
      (5): Linear(in_features=256, out_features=256, bias=False)
      (6): Dropout(p=0, inplace=False)
    )
    (1-2): 2 x ModuleList(
      (0): Linear(in_features=256, out_features=256, bias=True)
      (1): Linear(in_features=256, out_features=256, bias=False)
      (2): Linear(in_features=256, out_features=256, bias=True)
      (3): Linear(in_features=256, out_features=256, bias=False)
      (4): Linear(in_features=256, out_features=256, bias=True)
      (5): Linear(in_features=256, out_features=256, bias=False)
      (6): Dropout(p=0, inplace=False)
    )
    (3): ModuleList(
      (0): Linear(in_features=256, out_features=78, bias=True)
    )
  )
  (out_layer): Linear(in_features=256, out_features=78, bias=True)
)

y.shape=torch.Size([32, 64, 78])
h.shape=torch.Size([32, 3, 256])

scores = y[0,0,:].detach()
_, ax = plt.subplots(figsize=(15,5))

for t in reversed([0.3, 0.5, 1.0, 100]):
    ax.plot(charnn.hot_softmax(scores, temperature=t).cpu().numpy(), label=f'T={t}')
ax.set_xlabel('$x_{t+1}$')
ax.set_ylabel('$p(x_{t+1}|x_t)$')
ax.legend()

uniform_proba = 1/len(char_to_idx)
uniform_diff = torch.abs(charnn.hot_softmax(scores, temperature=100) - uniform_proba)
test.assertTrue(torch.all(uniform_diff < 1e-4))

for _ in range(3):
    text = charnn.generate_from_model(model, "foobar", 50, (char_to_idx, idx_to_char), T=0.5)
    print(text)
    test.assertEqual(len(text), 50)

foobarh0d,hBnOiN)!RyA3qLn0okBsDDhVF?7z
p4:GrTTIYe8
foobartZ:t]zsCeM6fdQ&?oJ7S93W!RC](vm6Y!-rl  8&gk]d
foobarJ!aMooL5PRYbV'HQ7r; j[Uv1fp[VroEpt(&Pm"mx?]4

# Pick a tiny subset of the dataset
subset_start, subset_end = 1001, 1005
ds_corpus_ss = torch.utils.data.Subset(ds_corpus, range(subset_start, subset_end))
batch_size_ss = 1
sampler_ss = SequenceBatchSampler(ds_corpus_ss, batch_size=batch_size_ss)
dl_corpus_ss = torch.utils.data.DataLoader(ds_corpus_ss, batch_size_ss, sampler=sampler_ss, shuffle=False)

# Convert subset to text
subset_text = ''
for i in range(subset_end - subset_start):
    subset_text += unembed(ds_corpus_ss[i][0])
print(f'Text to "memorize":\n\n{subset_text}')

Text to "memorize":

TRAM. What would you have?
  HELENA. Something; and scarce so much; nothing, indeed.
    I would not tell you what I would, my lord.
    Faith, yes:
    Strangers and foes do sunder and not kiss.
  BERTRAM. I pray you, stay not, but in haste to horse.
  HE

import torch.nn as nn
import torch.optim as optim
from hw3.training import RNNTrainer

torch.manual_seed(42)

lr = 0.01
num_epochs = 500

in_dim = vocab_len
h_dim = 128
n_layers = 2
loss_fn = nn.CrossEntropyLoss()
model = charnn.MultilayerGRU(in_dim, h_dim, out_dim=in_dim, n_layers=n_layers).to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
trainer = RNNTrainer(model, loss_fn, optimizer, device)

for epoch in range(num_epochs):
    epoch_result = trainer.train_epoch(dl_corpus_ss, verbose=False)
    
    # Every X epochs, we'll generate a sequence starting from the first char in the first sequence
    # to visualize how/if/what the model is learning.
    if epoch == 0 or (epoch+1) % 25 == 0:
        avg_loss = np.mean(epoch_result.losses)
        accuracy = np.mean(epoch_result.accuracy)
        print(f'\nEpoch #{epoch+1}: Avg. loss = {avg_loss:.3f}, Accuracy = {accuracy:.2f}%')
        
        generated_sequence = charnn.generate_from_model(model, subset_text[0],
                                                        seq_len*(subset_end-subset_start),
                                                        (char_to_idx,idx_to_char), T=0.1)
        
        # Stop if we've successfully memorized the small dataset.
        print(generated_sequence)
        if generated_sequence == subset_text:
            break

# Test successful overfitting
test.assertGreater(epoch_result.accuracy, 99)
test.assertEqual(generated_sequence, subset_text)

Epoch #1: Avg. loss = 3.937, Accuracy = 17.97%

Ttn                         t            t     t  n t                       t    t         t                                   t       t    tt           t       t                             t    t                          t           t          t

Epoch #25: Avg. loss = 1.028, Accuracy = 73.44%

TAAM. What would you hat youhde?
  HELLLAAA. Sotethang; and soat yould not teas andel you teinge inde, note ind soat nothinge indee and sou teas andes and yoat sould not tel, notete so yoas indeesangete soet nothinge indee inde, note so yoas indeesand you

Epoch #50: Avg. loss = 0.027, Accuracy = 100.00%

TRAM. What would you have?
  HELENA. Something; and scarce so much; nothing, indeed.
    I would not tell you what I would, my lord.
    Faith, yes:
    Strangers and foes do sunder and not kiss.
  BERTRAM. I pray you, stay not, but in haste to horse.
  HE

from hw3.answers import part1_rnn_hyperparams

hp = part1_rnn_hyperparams()
print('hyperparams:\n', hp)

### Dataset definition
vocab_len = len(char_to_idx)
batch_size = hp['batch_size']
seq_len = hp['seq_len']
train_test_ratio = 0.9
num_samples = (len(corpus) - 1) // seq_len
num_train = int(train_test_ratio * num_samples)

samples, labels = charnn.chars_to_labelled_samples(corpus, char_to_idx, seq_len, device)

ds_train = torch.utils.data.TensorDataset(samples[:num_train], labels[:num_train])
sampler_train = SequenceBatchSampler(ds_train, batch_size)
dl_train = torch.utils.data.DataLoader(ds_train, batch_size, shuffle=False, sampler=sampler_train, drop_last=True)

ds_test = torch.utils.data.TensorDataset(samples[num_train:], labels[num_train:])
sampler_test = SequenceBatchSampler(ds_test, batch_size)
dl_test = torch.utils.data.DataLoader(ds_test, batch_size, shuffle=False, sampler=sampler_test, drop_last=True)

print(f'Train: {len(dl_train):3d} batches, {len(dl_train)*batch_size*seq_len:7d} chars')
print(f'Test:  {len(dl_test):3d} batches, {len(dl_test)*batch_size*seq_len:7d} chars')

### Training definition
in_dim = out_dim = vocab_len
checkpoint_file = 'checkpoints/rnn'
num_epochs = 50
early_stopping = 5

model = charnn.MultilayerGRU(in_dim, hp['h_dim'], out_dim, hp['n_layers'], hp['dropout'])
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=hp['learn_rate'])
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='max', factor=hp['lr_sched_factor'], patience=hp['lr_sched_patience'], verbose=True
)
trainer = RNNTrainer(model, loss_fn, optimizer, device)

hyperparams:
 {'batch_size': 64, 'seq_len': 64, 'h_dim': 256, 'n_layers': 3, 'dropout': 0.2, 'learn_rate': 0.001, 'lr_sched_factor': 0.1, 'lr_sched_patience': 2}

Train: 1394 batches, 5709824 chars
Test:  154 batches,  630784 chars

/home/eran.b/miniconda3/envs/cs236781-hw/lib/python3.8/site-packages/torch/utils/data/sampler.py:65: UserWarning: `data_source` argument is not used and will be removed in 2.2.0.You may still have custom implementation that utilizes it.
  warnings.warn("`data_source` argument is not used and will be removed in 2.2.0."

from cs236781.plot import plot_fit

def post_epoch_fn(epoch, train_res, test_res, verbose):
    # Update learning rate
    scheduler.step(test_res.accuracy)
    # Sample from model to show progress
    if verbose:
        start_seq = "ACT I."
        generated_sequence = charnn.generate_from_model(
            model, start_seq, 100, (char_to_idx,idx_to_char), T=0.5
        )
        print(generated_sequence)

# Train, unless final checkpoint is found
checkpoint_file_final = f'{checkpoint_file}_final.pt'
if os.path.isfile(checkpoint_file_final):
    print(f'*** Loading final checkpoint file {checkpoint_file_final} instead of training')
    saved_state = torch.load(checkpoint_file_final, map_location=device)
    model.load_state_dict(saved_state['model_state'])
else:
    try:
        # Print pre-training sampling
        print(charnn.generate_from_model(model, "ACT I.", 100, (char_to_idx,idx_to_char), T=0.5))
    
        fit_res = trainer.fit(dl_train, dl_test, num_epochs, max_batches=None,
                                  post_epoch_fn=post_epoch_fn, early_stopping=early_stopping,
                                  checkpoints=checkpoint_file, print_every=1)
            
        fig, axes = plot_fit(fit_res)
    except KeyboardInterrupt as e:
            print('\n *** Training interrupted by user')

/home/eran.b/HW3/hw3/training.py:76: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  saved_state = torch.load(checkpoint_filename, map_location=self.device)

ACT I.xgqb[nwv.EfEQvPeoTxX,.l Mvh0AWXFvZkBs58g"ppO?EOs2com&w4E1?:0bBzbH;aMJj;b qBpVZ3n-YT."&Xizh
C)V
*** Loading checkpoint file checkpoints/rnn.pt
--- EPOCH 1/50 ---

train_batch:   0%|          | 0/1394 [00:00<?, ?it/s]

train_batch (1.088):   0%|          | 0/1394 [00:00<?, ?it/s]

train_batch (1.088):   0%|          | 1/1394 [00:00<07:02,  3.29it/s]

train_batch (0.991):   0%|          | 1/1394 [00:00<07:02,  3.29it/s]

train_batch (0.991):   0%|          | 2/1394 [00:00<06:41,  3.46it/s]

train_batch (1.036):   0%|          | 2/1394 [00:00<06:41,  3.46it/s]

train_batch (1.036):   0%|          | 3/1394 [00:00<06:34,  3.53it/s]

train_batch (1.000):   0%|          | 3/1394 [00:01<06:34,  3.53it/s]

from hw3.answers import part1_generation_params

start_seq, temperature = part1_generation_params()

generated_sequence = charnn.generate_from_model(
    model, start_seq, 10000, (char_to_idx,idx_to_char), T=temperature
)

print(generated_sequence)

ACT I.
                      Alexander. Five babes are gallows,
             What he is but this bed,
                  His friends will formal,
                      And hold the stoner]  
                            A castle.

    How now? Who rais'd thee?  
  CAESAR. This is a rack accompant.
  TITUS. Nay, gentle Ere die with my spirits on.
    Sound Caesar, go to Caesar, if thou hast
    To make his forehead or such a giberous lord.
  ANNE. Your sword counsel, have I tonight.
  THIRD CONSPIRATOR. He shall not be so five to the princes.
  FIRST WATCH. Provoked a gentleman.
                                                           [Kisses him]
    And therefore should be contented to him.
  PANDULPH. He is our passion, sweet light or two words,
    Put them home where all, flattery, and free and slave!
    My means to my soul to his life, your hateful doors,
    Doing the pride of Cade, corely come to you.
    I come to her despite. All the flowers
    Whose tender arms is angry with our meaning.
    Shall I be crown'd, and stay not with his side.
    This might have the sport a wingeth council,
    King Henry gives me no man doth mine,
    Can make them bring to the wretched man.
  BASSANIO. You shall please you, anon, and let this charge;
    Some confederacy need not see the King.
  TRANIO. I do not speak but that he is a maid,
    Thou wilt fall that duke nor of the King,
    And with her sons was well due in power
    When an 'Gius is, it seems hard by his professed
    As to point my conscience. Now I see it,
    That I could have been proud. Farewell to you.
  CHARMIAN. I do not know her to their beds; he does be.
    And for these prejugities were all of all
    For this sword and to do with me with his bragghor
    And not my sons- for the other sands,
    Nor from this few monusts, 'for his knife!
    What will I, more than a port as an
    Scalits to the warm Sourching upon this babe,
    Where been as sending all of sleep,
    Fall occation or prey.
  DUKE OF VENICE. Why, good my lord,
    Where's it entreat his power? Do you think,
    Or was she return'd by the customary  
    That he hath a main terms, the interumory,
    That deposed sorrows of the face of hell,
    And wide death broke thine ears and actions
    With thine evil, from whence you will make the vall'd
    The great sound which bears me, should have themselves,
    For some of thine eyes and every second
    Worthy of infections. Hector, we have mounted
    More of the common ground would conceive
    That every stroke of salt with weapons will,
    Come then, being what shall charge the deed,
    Should so corrupt against thy books upon our
    thanks and that.
  SPEED. 'Tis his speech were owed, lest you shall with me down.
  ORLANDO. Do you not thank she? My mother, have you undergo me
    upon him? Well, heaven nor hither we are so having to-day, I will say,
    friend or a fool to be too run'd; there are no such friends! How say you,
    that any knave now to make my play to bed of the throng, or
    he is a good action! Whose water'd weak as that, in mine has
    filing his master.
  SECOND DOCKS. Will he come? What will you?
  BERTRAM. O, the door, good friends. I pray you, let him know
    She was more career, and the law of Cymbeline
    I swear, fresh rebouts for mercy. If 'twere so,
    Each men of which he gave him in the flower
    That Annidus wings, flattering affections!
    Virtue is chider. But all the manner of him.
  LUCILIUS. Great that Julius Caesar weep or affection.
    I do inflier a crest on the seas  
    Are mischief, the truth of my forehead,
    And there abholusion and most hardinant
    I should not speak.
    Thy health, I ask all foel of fear.
    I cannot swear high love as dearer than
    And say 'God, I am content.' 'Scour we pay?
    I serve thee well; but he is very much.
  KING EDWARD. Say 'tis thought of.
  LUCIUS.  [Aside]  The first saw the very friend, yea,
    Upon the bloody deed of nature, now
    I will speak a ludgy unworthy a monther.
    One simple to thy head among the blest,
    If we with your mouth to arms, and other
    Even as she is not thy heavy hand,
    And notice, if she be found mine eyes,
    Should be off to speak- I am ta'en a man.
    Welcome hards, and have him prepare thee up.
    This drum to me, where hast thou been too sure,
    He would be the first to dash the green helms?
    Come, master, give me this night confirmite.
    I am charg'd. There is no bodies will be most
    He's levy.
  CORIOLANUS. Hear me.                                Exeunt

                         SCENE II.
                                     Rome his son
                               Still flinty enter Roderigo.
  LEPIDUS. Go, publicate that I bring him breath;
    And now they shall have good or net but a
    legs, and I pray, though I shall pray for you. But give it with him.
  ALL. Why should he repent?
  TOUCHSTONE. He longs for a woman. He is worse to change four easing that I
    swallowed with the charge.
    It was such petty serious. I am your offer'd by him. Come,
    what would you have you die? And I fear the matter; and what's
    a term for thee again with a knight?
  ISABELLA. I have nobody, be it so, my liege,
    Camillo with fair and worse forth that shall,
    With thy lord, sir, then why thou cross 'a ground;
    And in that fellow shall chance to hold
    The love-sovereign fellow seem so.  
                                                    [KANGERNs] Grace is from the Romans.
    Are you here expire your English judge to you?
  VIOLA. Good morrow, man; there is a trick of spirits
    To any fair shame with a little sun
    Is enter-company. Let me seen a gentleman,
    Lest I can tell you both a promise.
    Why dost thou find out this maid nor lose,
    The rest I think he did conclude for state?
    Poor were eleven good as he reads.
    Go, silence. This summer of the city  
    Is not a husband beseech you, lest they do
    With childishry as thou art all the court,
    To fear his grace, but stealing on my life
    Before and natural sound of the trial.
    Cry 'Ophor's death and commander,
    Bear's his daughter's meaner blood, that makes
    Here to proclaim our standard. There thou less
    Become the fashion of that hand on thee,
    And for an effected in the head,
    Seems or could have better peril of my soul.
    Bade your pains and what I throw away how heaven!
    I beseech your Grace to change my staff,
    And when the picture is as a man;
    And if it be that ancient pleasure,
    Then will I give thee shame to see the comfort,
    Upon what exchange lov'd her own reason,
    By this melancholy vain concealment.
    So may you bury him?
  LEPIDUS. What's this?
    You are like to Macbeth, her friendly spirits-
    The prisoner should not appear upon us.
    That is ready, and pass we should be thee.
  PETRUCHIO. Upon my knee would try you out at Bassanio,
    Call thy subjects do chance to fight again,
    The house, the King, both walking cloaks of me.
    My Lord of York lays the pow'r struck or shore,
    To the best things will show thee from the island
    That drop everen do make him grows the air,
    And the rich manly devils bear down the hand,  
    And the action of a fury of her urge
    The heavy cause can seek the bargain I
    Into your brother's powers that thou shalt bear
    Upon their life, may must anon. Farewell, farewell.
    But ill, sir, we ure no more of the tops.
  PETRUCHIO. A persuasion! Sirrah, what says your mind?  
    Make no matter to me in courtesy
    Were then to our comfort. And thou, Lord Lords,
    My mother sits his fortune. Come, let's requite
    A notard; but she hath been to heaven;
    And as if many cur of sorrow liv'd,
    Shall be to rebel the rest, and three horses
    Of cruel argument, and my hounds!
  GLOUCESTER. She should not speak 'main so love-swearing
    To obey the other that I receiv'd
    The blood o' th' law-pompes of perfect in thin'd.
    Come, man! What men of sweating? What's this?
    He's a neigh with you. I have a purpose
    Have made thee cold.
  MENENIUS. Why, how, ho!
  FIRST LORD. You take my heart that wonders too; I do dispose
    His horns most degree all perpetual marvelle.
    At old Sir Richard Ratcliff. Look, we'll presume
    To leave from our husbandry; yet I shall,
    And in this many may be sounded thee.
    Thou canst not be there, and first, and toward my life,
    Little to our heads of fear and being silence.
  KING PHILIP. Ay, that in an honourable gentleman
    Have bide the nature of your foes as he.
  PORTIA. I cry you mercy, to my name will not.
  CATESBY. I come, Julia.
  MARIA. 'Tis not what I have made us'd a man
    To glorious ear into the belly.
    This is the course of stately gilded gird.
  FIRST SOLDIER. Sweet Sir Neaple, with tears do not put my words.
    There is no need well shalt search me in this heak.
    What profane morning lacks thence away with speed?
                                                              Exeunt

SCENE IV.
Another part of them in Scythian, and Attendants.

  HERALD. I have been small husband and your eye in England
    That she is not to partner.
  TIMON. Excellent peace.
    But by my trespass I do not live.
    You must parley.
  LEONTES. What! away with me?
  IMOGEN. I remember this;
    And then the merry food of courtier.
  GLOUCESTER. The ancient grace shall devour this new yet.
  BUCKINGHAM. Well, give me joy recounters his steward.
    To the Lady, didst thou have thy soul to life.
    What would he confess, for the next youth and,
    Thy vassal made twenty times to see,
    And show this new-made pollute wild offers;
    Patience to our cause alone, and fair
    On him, my mind so bells, the benefit
    Like a frosty kindred teams her chariot,
    And wash his rich siege. But why, she is here!-
    The like untainted is too soft accord
    To see the sea-mouse of proud to defend.
    The

from cs236781.answers import display_answer
import hw3.answers

display_answer(hw3.answers.part1_q1)

display_answer(hw3.answers.part1_q2)

display_answer(hw3.answers.part1_q3)

display_answer(hw3.answers.part1_q4)

import unittest
import os
import sys
import pathlib
import urllib
import shutil
import re
import zipfile

import numpy as np
import torch
import matplotlib.pyplot as plt

%load_ext autoreload
%autoreload 2

test = unittest.TestCase()
plt.rcParams.update({'font.size': 12})
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)

Using device: cuda

import cs236781.plot as plot
import cs236781.download
from hw3.answers import PART2_CUSTOM_DATA_URL as CUSTOM_DATA_URL

DATA_DIR = pathlib.Path.home().joinpath('.pytorch-datasets')
if CUSTOM_DATA_URL is None:
    DATA_URL = 'https://github.com/AviaAvraham1/TempDatasets/raw/refs/heads/main/George_W_Bush2.zip'
else:
    DATA_URL = CUSTOM_DATA_URL

_, dataset_dir = cs236781.download.download_data(out_path=DATA_DIR, url=DATA_URL, extract=True, force=False)

Downloading https://github.com/AviaAvraham1/TempDatasets/raw/refs/heads/main/George_W_Bush2.zip...

Saved to /home/eran.b/.pytorch-datasets/George_W_Bush2.zip.
Extracting /home/eran.b/.pytorch-datasets/George_W_Bush2.zip...

Extracted 532 to /home/eran.b/.pytorch-datasets/George_W_Bush/George_W_Bush

import torchvision.transforms as T
from torchvision.datasets import ImageFolder

im_size = 64
tf = T.Compose([
    # Resize to constant spatial dimensions
    T.Resize((im_size, im_size)),
    # PIL.Image -> torch.Tensor
    T.ToTensor(),
    # Dynamic range [0,1] -> [-1, 1]
    T.Normalize(mean=(.5,.5,.5), std=(.5,.5,.5)),
])

ds_gwb = ImageFolder(os.path.dirname(dataset_dir), tf)

_ = plot.dataset_first_n(ds_gwb, 50, figsize=(15,10), nrows=5)
print(f'Found {len(ds_gwb)} images in dataset folder.')

Found 530 images in dataset folder.

x0, y0 = ds_gwb[0]
x0 = x0.unsqueeze(0).to(device)
print(x0.shape)

test.assertSequenceEqual(x0.shape, (1, 3, im_size, im_size))

torch.Size([1, 3, 64, 64])

import hw3.autoencoder as autoencoder

in_channels = 3
out_channels = 1024
encoder_cnn = autoencoder.EncoderCNN(in_channels, out_channels).to(device)
print(encoder_cnn)

h = encoder_cnn(x0)
print(h.shape)

test.assertEqual(h.dim(), 4)
test.assertSequenceEqual(h.shape[0:2], (1, out_channels))

EncoderCNN(
  (cnn): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
    (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): ReLU(inplace=True)
    (6): Conv2d(128, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (7): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (8): ReLU(inplace=True)
  )
)

torch.Size([1, 1024, 8, 8])

decoder_cnn = autoencoder.DecoderCNN(in_channels=out_channels, out_channels=in_channels).to(device)
print(decoder_cnn)
x0r = decoder_cnn(h)
print(x0r.shape)

test.assertEqual(x0.shape, x0r.shape)

# Should look like colored noise
T.functional.to_pil_image(x0r[0].cpu().detach())

DecoderCNN(
  (cnn): Sequential(
    (0): ConvTranspose2d(1024, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)
    (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
    (3): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)
    (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): ReLU(inplace=True)
    (6): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))
  )
)

torch.Size([1, 3, 64, 64])

z_dim = 2
vae = autoencoder.VAE(encoder_cnn, decoder_cnn, x0[0].size(), z_dim).to(device)
print(vae)

z, mu, log_sigma2 = vae.encode(x0)

test.assertSequenceEqual(z.shape, (1, z_dim))
test.assertTrue(z.shape == mu.shape == log_sigma2.shape)

print(f'mu(x0)={list(*mu.detach().cpu().numpy())}, sigma2(x0)={list(*torch.exp(log_sigma2).detach().cpu().numpy())}')

VAE(
  (features_encoder): EncoderCNN(
    (cnn): Sequential(
      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (5): ReLU(inplace=True)
      (6): Conv2d(128, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (7): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (8): ReLU(inplace=True)
    )
  )
  (features_decoder): DecoderCNN(
    (cnn): Sequential(
      (0): ConvTranspose2d(1024, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)
      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)
      (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (5): ReLU(inplace=True)
      (6): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))
    )
  )
  (project_mu): Linear(in_features=65536, out_features=2, bias=True)
  (project_logvar): Linear(in_features=65536, out_features=2, bias=True)
  (project_z): Linear(in_features=2, out_features=65536, bias=True)
)

mu(x0)=[0.8224494, -0.51938593], sigma2(x0)=[0.86267585, 1.9702746]

# Sample from q(Z|x)
N = 500
Z = torch.zeros(N, z_dim)
_, ax = plt.subplots()
with torch.no_grad():
    for i in range(N):
        Z[i], _, _ = vae.encode(x0)
        ax.scatter(*Z[i].cpu().numpy())

# Should be close to the mu/sigma in the previous block above
print('sampled mu', torch.mean(Z, dim=0))
print('sampled sigma2', torch.var(Z, dim=0))

sampled mu tensor([ 0.8101, -0.4607])
sampled sigma2 tensor([0.9138, 1.8869])

x0r = vae.decode(z)

test.assertSequenceEqual(x0r.shape, x0.shape)

x0r, mu, log_sigma2 = vae(x0)

test.assertSequenceEqual(x0r.shape, x0.shape)
test.assertSequenceEqual(mu.shape, (1, z_dim))
test.assertSequenceEqual(log_sigma2.shape, (1, z_dim))
T.functional.to_pil_image(x0r[0].detach().cpu())

from hw3.autoencoder import vae_loss
torch.manual_seed(42)

def test_vae_loss():
    # Test data
    N, C, H, W = 10, 3, 64, 64 
    z_dim = 32
    x  = torch.randn(N, C, H, W)*2 - 1
    xr = torch.randn(N, C, H, W)*2 - 1
    z_mu = torch.randn(N, z_dim)
    z_log_sigma2 = torch.randn(N, z_dim)
    x_sigma2 = 0.9
    
    loss, _, _ = vae_loss(x, xr, z_mu, z_log_sigma2, x_sigma2)
    
    test.assertAlmostEqual(loss.item(), 58.3234367, delta=1e-3)
    return loss

test_vae_loss()

tensor(58.3234)

samples = vae.sample(5)
_ = plot.tensors_as_images(samples)

import torch.optim as optim
from torch.utils.data import random_split
from torch.utils.data import DataLoader
from torch.nn import DataParallel
from hw3.training import VAETrainer
from hw3.answers import part2_vae_hyperparams

torch.manual_seed(42)

# Hyperparams
hp = part2_vae_hyperparams()
batch_size = hp['batch_size']
h_dim = hp['h_dim']
z_dim = hp['z_dim']
x_sigma2 = hp['x_sigma2']
learn_rate = hp['learn_rate']
betas = hp['betas']

# Data
split_lengths = [int(len(ds_gwb)*0.9), int(len(ds_gwb)*0.1)]
ds_train, ds_test = random_split(ds_gwb, split_lengths)
dl_train = DataLoader(ds_train, batch_size, shuffle=True)
dl_test  = DataLoader(ds_test,  batch_size, shuffle=True)
im_size = ds_train[0][0].shape

# Model
encoder = autoencoder.EncoderCNN(in_channels=im_size[0], out_channels=h_dim)
decoder = autoencoder.DecoderCNN(in_channels=h_dim, out_channels=im_size[0])
vae = autoencoder.VAE(encoder, decoder, im_size, z_dim)
vae_dp = DataParallel(vae).to(device)

# Optimizer
optimizer = optim.Adam(vae.parameters(), lr=learn_rate, betas=betas)

# Loss
def loss_fn(x, xr, z_mu, z_log_sigma2):
    return autoencoder.vae_loss(x, xr, z_mu, z_log_sigma2, x_sigma2)

# Trainer
trainer = VAETrainer(vae_dp, loss_fn, optimizer, device)
checkpoint_file = 'checkpoints/vae'
checkpoint_file_final = f'{checkpoint_file}_final'
if os.path.isfile(f'{checkpoint_file}.pt'):
    os.remove(f'{checkpoint_file}.pt')

# Show model and hypers
print(vae)
print(hp)

VAE(
  (features_encoder): EncoderCNN(
    (cnn): Sequential(
      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (5): ReLU(inplace=True)
      (6): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (7): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (8): ReLU(inplace=True)
    )
  )
  (features_decoder): DecoderCNN(
    (cnn): Sequential(
      (0): ConvTranspose2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)
      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)
      (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (5): ReLU(inplace=True)
      (6): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))
    )
  )
  (project_mu): Linear(in_features=8192, out_features=128, bias=True)
  (project_logvar): Linear(in_features=8192, out_features=128, bias=True)
  (project_z): Linear(in_features=128, out_features=8192, bias=True)
)
{'batch_size': 32, 'h_dim': 128, 'z_dim': 128, 'x_sigma2': 0.9, 'learn_rate': 0.0002, 'betas': (0.5, 0.999)}

import IPython.display

def post_epoch_fn(epoch, train_result, test_result, verbose):
    # Plot some samples if this is a verbose epoch
    if verbose:
        samples = vae.sample(n=5)
        fig, _ = plot.tensors_as_images(samples, figsize=(6,2))
        IPython.display.display(fig)
        plt.close(fig)

if os.path.isfile(f'{checkpoint_file_final}.pt'):
    print(f'*** Loading final checkpoint file {checkpoint_file_final} instead of training')
    checkpoint_file = checkpoint_file_final
else:
    res = trainer.fit(dl_train, dl_test,
                      num_epochs=200, early_stopping=20, print_every=10,
                      checkpoints=checkpoint_file,
                      post_epoch_fn=post_epoch_fn)
    
# Plot images from best model
saved_state = torch.load(f'{checkpoint_file}.pt', map_location=device)
vae_dp.load_state_dict(saved_state['model_state'])
print('*** Images Generated from best model:')
fig, _ = plot.tensors_as_images(vae_dp.module.sample(n=15), nrows=3, figsize=(6,6))

--- EPOCH 1/200 ---

train_batch:   0%|          | 0/15 [00:00<?, ?it/s]

train_batch (33.956):   0%|          | 0/15 [00:01<?, ?it/s]

train_batch (33.956):   7%|▋         | 1/15 [00:01<00:14,  1.07s/it]

train_batch (46.432):   7%|▋         | 1/15 [00:01<00:14,  1.07s/it]

train_batch (22.703):  13%|█▎        | 2/15 [00:01<00:13,  1.07s/it]

train_batch (22.703):  20%|██        | 3/15 [00:01<00:04,  2.99it/s]

train_batch (23.977):  20%|██        | 3/15 [00:01<00:04,  2.99it/s]

train_batch (18.520):  27%|██▋       | 4/15 [00:01<00:03,  2.99it/s]

train_batch (18.520):  33%|███▎      | 5/15 [00:01<00:02,  4.98it/s]

from cs236781.answers import display_answer
import hw3.answers

display_answer(hw3.answers.part2_q1)

display_answer(hw3.answers.part2_q2)

display_answer(hw3.answers.part2_q3)

display_answer(hw3.answers.part2_q4)

%load_ext autoreload
%autoreload 2
import unittest
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import torch.optim as optim
from tqdm import tqdm
import os
import numpy as np

test = unittest.TestCase()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print('Using device:', device)

Using device: cuda

from hw3.transformer import sliding_window_attention


## test sliding-window attention
num_heads = 3
batch_size = 2
seq_len = 5
embed_dim = 3
window_size = 2

## test without extra dimension for heads
x = torch.arange(seq_len*embed_dim).reshape(seq_len,embed_dim).repeat(batch_size,1).reshape(batch_size, seq_len, -1).float()

values, attention = sliding_window_attention(x, x, x,window_size)

gt_values = torch.load(os.path.join('test_tensors','values_tensor_0_heads.pt'))


test.assertTrue(torch.all(values == gt_values), f'the tensors differ in dims [B,row,col]:{torch.stack(torch.where(values != gt_values),dim=0)}')

gt_attention = torch.load(os.path.join('test_tensors','attention_tensor_0_heads.pt'))
test.assertTrue(torch.all(attention == gt_attention), f'the tensors differ in dims [B,row,col]:{torch.stack(torch.where(attention != gt_attention),dim=0)}')


## test with extra dimension for heads
x = torch.arange(seq_len*embed_dim).reshape(seq_len,embed_dim).repeat(batch_size, num_heads, 1).reshape(batch_size, num_heads, seq_len, -1).float()

values, attention = sliding_window_attention(x, x, x,window_size)

gt_values = torch.load(os.path.join('test_tensors','values_tensor_3_heads.pt'))
test.assertTrue(torch.all(values == gt_values), f'the tensors differ in dims [B,num_heads,row,col]:{torch.stack(torch.where(values != gt_values),dim=0)}')


gt_attention = torch.load(os.path.join('test_tensors','attention_tensor_3_heads.pt'))
test.assertTrue(torch.all(attention == gt_attention), f'the tensors differ in dims [B,num_heads,row,col]:{torch.stack(torch.where(attention != gt_attention),dim=0)}')

/tmp/ipykernel_1886718/1539104548.py:16: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  gt_values = torch.load(os.path.join('test_tensors','values_tensor_0_heads.pt'))
/tmp/ipykernel_1886718/1539104548.py:21: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  gt_attention = torch.load(os.path.join('test_tensors','attention_tensor_0_heads.pt'))
/tmp/ipykernel_1886718/1539104548.py:30: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  gt_values = torch.load(os.path.join('test_tensors','values_tensor_3_heads.pt'))
/tmp/ipykernel_1886718/1539104548.py:34: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  gt_attention = torch.load(os.path.join('test_tensors','attention_tensor_3_heads.pt'))

import numpy as np
import pandas as pd
import sys
import pathlib
import urllib
import shutil
import re

import matplotlib.pyplot as plt

%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

from datasets import DatasetDict
from datasets import load_dataset

/home/eran.b/miniconda3/envs/cs236781-hw/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

dataset = load_dataset('imdb', split=['train', 'test', 'train[12480:12520]'])

print(dataset)

[Dataset({
    features: ['text', 'label'],
    num_rows: 25000
}), Dataset({
    features: ['text', 'label'],
    num_rows: 25000
}), Dataset({
    features: ['text', 'label'],
    num_rows: 40
})]

#wrap it in a DatasetDict to enable methods such as map and format
dataset = DatasetDict({'train': dataset[0], 'val': dataset[1], 'toy': dataset[2]})

dataset

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    val: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    toy: Dataset({
        features: ['text', 'label'],
        num_rows: 40
    })
})

print(dataset['train'])

for i in range(4):
    print(f'TRAINING SAMPLE {i}:') 
    print(dataset['train'][i]['text'])
    label = dataset['train'][i]['label']
    print(f'Label {i}: {label}')
    print('\n')

Dataset({
    features: ['text', 'label'],
    num_rows: 25000
})
TRAINING SAMPLE 0:
I rented I AM CURIOUS-YELLOW from my video store because of all the controversy that surrounded it when it was first released in 1967. I also heard that at first it was seized by U.S. customs if it ever tried to enter this country, therefore being a fan of films considered "controversial" I really had to see this for myself.<br /><br />The plot is centered around a young Swedish drama student named Lena who wants to learn everything she can about life. In particular she wants to focus her attentions to making some sort of documentary on what the average Swede thought about certain political issues such as the Vietnam War and race issues in the United States. In between asking politicians and ordinary denizens of Stockholm about their opinions on politics, she has sex with her drama teacher, classmates, and married men.<br /><br />What kills me about I AM CURIOUS-YELLOW is that 40 years ago, this was considered pornographic. Really, the sex and nudity scenes are few and far between, even then it's not shot like some cheaply made porno. While my countrymen mind find it shocking, in reality sex and nudity are a major staple in Swedish cinema. Even Ingmar Bergman, arguably their answer to good old boy John Ford, had sex scenes in his films.<br /><br />I do commend the filmmakers for the fact that any sex shown in the film is shown for artistic purposes rather than just to shock people and make money to be shown in pornographic theaters in America. I AM CURIOUS-YELLOW is a good film for anyone wanting to study the meat and potatoes (no pun intended) of Swedish cinema. But really, this film doesn't have much of a plot.
Label 0: 0


TRAINING SAMPLE 1:
"I Am Curious: Yellow" is a risible and pretentious steaming pile. It doesn't matter what one's political views are because this film can hardly be taken seriously on any level. As for the claim that frontal male nudity is an automatic NC-17, that isn't true. I've seen R-rated films with male nudity. Granted, they only offer some fleeting views, but where are the R-rated films with gaping vulvas and flapping labia? Nowhere, because they don't exist. The same goes for those crappy cable shows: schlongs swinging in the breeze but not a clitoris in sight. And those pretentious indie movies like The Brown Bunny, in which we're treated to the site of Vincent Gallo's throbbing johnson, but not a trace of pink visible on Chloe Sevigny. Before crying (or implying) "double-standard" in matters of nudity, the mentally obtuse should take into account one unavoidably obvious anatomical difference between men and women: there are no genitals on display when actresses appears nude, and the same cannot be said for a man. In fact, you generally won't see female genitals in an American film in anything short of porn or explicit erotica. This alleged double-standard is less a double standard than an admittedly depressing ability to come to terms culturally with the insides of women's bodies.
Label 1: 0


TRAINING SAMPLE 2:
If only to avoid making this type of film in the future. This film is interesting as an experiment but tells no cogent story.<br /><br />One might feel virtuous for sitting thru it because it touches on so many IMPORTANT issues but it does so without any discernable motive. The viewer comes away with no new perspectives (unless one comes up with one while one's mind wanders, as it will invariably do during this pointless film).<br /><br />One might better spend one's time staring out a window at a tree growing.<br /><br />
Label 2: 0


TRAINING SAMPLE 3:
This film was probably inspired by Godard's Masculin, féminin and I urge you to see that film instead.<br /><br />The film has two strong elements and those are, (1) the realistic acting (2) the impressive, undeservedly good, photo. Apart from that, what strikes me most is the endless stream of silliness. Lena Nyman has to be most annoying actress in the world. She acts so stupid and with all the nudity in this film,...it's unattractive. Comparing to Godard's film, intellectuality has been replaced with stupidity. Without going too far on this subject, I would say that follows from the difference in ideals between the French and the Swedish society.<br /><br />A movie of its time, and place. 2/10.
Label 3: 0

def label_cnt(type):
    ds = dataset[type]
    size = len(ds)
    cnt= 0 
    for smp in ds:
        cnt += smp['label']
    print(f'negative samples in {type} dataset: {size - cnt}')
    print(f'positive samples in {type} dataset: {cnt}')
    
label_cnt('train')
label_cnt('val')
label_cnt('toy')

negative samples in train dataset: 12500
positive samples in train dataset: 12500

negative samples in val dataset: 12500
positive samples in val dataset: 12500
negative samples in toy dataset: 20
positive samples in toy dataset: 20

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
print("Tokenizer input max length:", tokenizer.model_max_length)
print("Tokenizer vocabulary size:", tokenizer.vocab_size)

Tokenizer input max length: 512
Tokenizer vocabulary size: 30522

def tokenize_text(batch):
    return tokenizer(batch["text"], truncation=True, padding=True)

def tokenize_dataset(dataset):
    dataset_tokenized = dataset.map(tokenize_text, batched=True, batch_size =None)
    return dataset_tokenized

dataset_tokenized = tokenize_dataset(dataset)

# we would like to work with pytorch so we can manually fine-tune
dataset_tokenized.set_format("torch", columns=["input_ids", "attention_mask", "label"])

# no need to parrarelize in this assignment
os.environ["TOKENIZERS_PARALLELISM"] = "false"

from torch.utils.data import DataLoader, Dataset

class IMDBDataset(Dataset):
    def __init__(self, dataset):
        self.ds = dataset

    def __getitem__(self, index):
        return self.ds[index]

    def __len__(self):
        return self.ds.num_rows

train_dataset = IMDBDataset(dataset_tokenized['train'])
val_dataset = IMDBDataset(dataset_tokenized['val'])
toy_dataset = IMDBDataset(dataset_tokenized['toy'])

dl_train,dl_val, dl_toy = [ 
    DataLoader(
    dataset=train_dataset,
    batch_size=12,
    shuffle=True, 
    num_workers=0
),
DataLoader(
    dataset=val_dataset,
    batch_size=12,
    shuffle=True, 
    num_workers=0
),
DataLoader(
    dataset=toy_dataset,
    batch_size=4,
    num_workers=0
)]

from hw3.transformer import EncoderLayer
# set torch seed for reproducibility
torch.manual_seed(0)
layer = EncoderLayer(embed_dim=16, hidden_dim=16, num_heads=4, window_size=4, dropout=0.1)

# load x and y
x = torch.load(os.path.join('test_tensors','encoder_layer_input.pt'))
y = torch.load(os.path.join('test_tensors','encoder_layer_output.pt'))
padding_mask = torch.ones(2, 10)
padding_mask[:, 5:] = 0

# forward pass
out = layer(x, padding_mask)
test.assertTrue(torch.allclose(out, y, atol=1e-6), 'output of encoder layer is incorrect')

/tmp/ipykernel_1886718/2496630226.py:7: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  x = torch.load(os.path.join('test_tensors','encoder_layer_input.pt'))
/tmp/ipykernel_1886718/2496630226.py:8: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  y = torch.load(os.path.join('test_tensors','encoder_layer_output.pt'))

tokenizer.convert_ids_to_tokens(dataset_tokenized['train'][0]['input_ids'])[:10]

['[CLS]', 'i', 'rented', 'i', 'am', 'curious', '-', 'yellow', 'from', 'my']

from hw3.transformer import Encoder

# set torch seed for reproducibility
torch.manual_seed(0)
encoder = Encoder(vocab_size=100, embed_dim=16, num_heads=4, num_layers=3, 
                  hidden_dim=16, max_seq_length=64, window_size=4, dropout=0.1)

# load x and y
x = torch.load(os.path.join('test_tensors','encoder_input.pt'))
y = torch.load(os.path.join('test_tensors','encoder_output.pt'))
#x = torch.randint(0, 100, (2, 64)).long()

padding_mask = torch.ones(2, 64)
padding_mask[:, 50:] = 0

# forward pass
out = encoder(x, padding_mask)
test.assertTrue(torch.allclose(out, y, atol=1e-6), 'output of encoder layer is incorrect')

/tmp/ipykernel_1886718/537645795.py:9: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  x = torch.load(os.path.join('test_tensors','encoder_input.pt'))
/tmp/ipykernel_1886718/537645795.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  y = torch.load(os.path.join('test_tensors','encoder_output.pt'))

from hw3.answers import part3_transformer_encoder_hyperparams

params = part3_transformer_encoder_hyperparams()
print(params)
embed_dim = params['embed_dim'] 
num_heads = params['num_heads']
num_layers = params['num_layers']
hidden_dim = params['hidden_dim']
window_size = params['window_size']
dropout = params['droupout']
lr = params['lr']

vocab_size = tokenizer.vocab_size
max_seq_length = tokenizer.model_max_length

max_batches_per_epoch = None
N_EPOCHS = 20

{'embed_dim': 128, 'num_heads': 4, 'num_layers': 2, 'hidden_dim': 128, 'window_size': 128, 'droupout': 0.0, 'lr': 0.001}

toy_model = Encoder(vocab_size, embed_dim, num_heads, num_layers, hidden_dim, max_seq_length, window_size, dropout=dropout).to(device)
toy_optimizer = optim.Adam(toy_model.parameters(), lr=lr)
criterion = nn.BCEWithLogitsLoss()

# fit your model
import pickle
if not os.path.exists('toy_transfomer_encoder.pt'):
    # overfit
    from hw3.training import TransformerEncoderTrainer
    toy_trainer = TransformerEncoderTrainer(toy_model, criterion, toy_optimizer, device=device)
    # set max batches per epoch
    _ = toy_trainer.fit(dl_toy, dl_toy, N_EPOCHS, checkpoints='toy_transfomer_encoder', max_batches=max_batches_per_epoch)

    

toy_saved_state = torch.load('toy_transfomer_encoder.pt')
toy_best_acc = toy_saved_state['best_acc']
toy_model.load_state_dict(toy_saved_state['model_state'])

/tmp/ipykernel_1886718/3834131296.py:12: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  toy_saved_state = torch.load('toy_transfomer_encoder.pt')

<All keys matched successfully>

test.assertTrue(toy_best_acc >= 95)

max_batches_per_epoch = 500
N_EPOCHS = 4

model = Encoder(vocab_size, embed_dim, num_heads, num_layers, hidden_dim, max_seq_length, window_size, dropout).to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)

# fit your model
import pickle
if not os.path.exists('trained_transfomer_encoder.pt'):
    from hw3.training import TransformerEncoderTrainer
    trainer = TransformerEncoderTrainer(model, criterion, optimizer, device=device)
    # set max batches per epoch
    _ = trainer.fit(dl_train, dl_val, N_EPOCHS, checkpoints='trained_transfomer_encoder', max_batches=max_batches_per_epoch)
    

saved_state = torch.load('trained_transfomer_encoder.pt')
best_acc = saved_state['best_acc']
model.load_state_dict(saved_state['model_state'])

/tmp/ipykernel_1886718/3339141018.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  saved_state = torch.load('trained_transfomer_encoder.pt')

<All keys matched successfully>

test.assertTrue(best_acc >= 65)

rand_index = torch.randint(len(dataset_tokenized['val']), (1,))
rand_index

tensor([19583])

sample = dataset['val'][rand_index]
sample['text']

['Best fan boy movie I\'ve ever watched save "Free Enterprise."<br /><br />In some ways it reminded me of an early Kevin Smith film.<br /><br />If you do any kind of role playing, this movie will likely have you laughing often at its insatiable fun. Don\'t expect a big budget here, the acting is also questionable at times, but it really adds to the fun of the atmosphere they create. The script is truly humorous with a lot of witty moments worth experiencing.<br /><br />The bard that always gets killed had me rolling. The sexually confused player also had me smiling a bit too. But in the end... It was just a great movie showcasing some better moments in the lives of a few geeks having a great time with role playing.<br /><br />If you are bored, and ever got into role playing, this will do nicely for a distraction. A real unexpected treat.']

tokenized_sample = dataset_tokenized['val'][rand_index]
tokenized_sample
input_ids = tokenized_sample['input_ids'].to(device)
label = tokenized_sample['label'].to(device)
attention_mask = tokenized_sample['attention_mask'].to(float).to(device)

print('label', label.shape)
print('attention_mask', attention_mask.shape)
prediction = model.predict(input_ids, attention_mask).squeeze(0)

print('label: {}, prediction: {}'.format(label, prediction))

label torch.Size([1])
attention_mask torch.Size([1, 512])

label: tensor([1], device='cuda:0'), prediction: tensor([1.], device='cuda:0', grad_fn=<SqueezeBackward1>)

from cs236781.answers import display_answer
import hw3.answers

display_answer(hw3.answers.part3_q1)

display_answer(hw3.answers.part3_q2)

#	Name	Id	email
Student 1	Ofir Fuchs	326141611	ofir.fuchs@campus.technion.ac.il
Student 2	Eran Ben Elihao	330596941	eran.b@campus.technion.ac.il

CS236781: Deep Learning on Computational Accelerators¶

Homework Assignment 3¶

Introduction¶

General Guidelines¶

Contents¶

Part 1: Sequence Models¶

Text generation with a char-level RNN¶

Obtaining the corpus¶

Data Preprocessing¶

Dataset Creation¶

Model Implementation¶

Generating text by sampling¶

Training¶

Generating a work of art¶

Questions¶

Question 1¶

Question 2¶

Question 3¶

Question 4¶

Part 2: Variational Autoencoder¶

Obtaining the dataset¶

The Variational Autoencoder¶

Model Implementation¶

Loss Implementation¶

Sampling¶

Training¶

Questions¶

Question 1¶

Question 2¶

Question 3¶

Question 4¶

Part 3: Transformer¶

Reminder: scaled dot product attention¶

Sliding window attention¶

Multihead Sliding window attention¶

Sentiment analysis¶

IMBD hugging face dataset¶

Import the tokenizer for the dataset¶

Setting up the dataloaders and dataset¶

Transformer Encoder¶

Training the encoder¶

Training on a toy dataset¶

Training on all data¶

Questions¶

Question 1¶

Question 2¶